import json
import asyncio
import logging
import paramiko
import threading
import time
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth import get_user_model
from django.utils import timezone
from .models import SSHSession, SSHCommand
from cmdb.models import Host
from common.security import (
    security_validator, ssh_rate_limiter, audit_logger,
    connection_pool
)

User = get_user_model()
logger = logging.getLogger(__name__)


class SSHTerminalConsumer(AsyncWebsocketConsumer):
    """SSH终端WebSocket消费者"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.session = None
        self.host = None
        self.user = None
        self.read_thread = None
        self.is_connected = False

    async def connect(self):
        """WebSocket连接"""
        # 获取用户信息
        self.user = self.scope["user"]
        if not self.user.is_authenticated:
            await self.close(code=4001)
            return

        # 速率限制检查
        client_ip = self.get_client_ip()
        user_identifier = f"{self.user.username}:{client_ip}"
        if not ssh_rate_limiter.is_allowed(user_identifier):
            logger.warning(f"SSH rate limit exceeded for user {self.user.username} from {client_ip}")
            await self.close(code=4029)  # Too Many Requests
            return

        # 获取主机ID
        self.host_id = self.scope['url_route']['kwargs']['host_id']

        try:
            self.host = await self.get_host(self.host_id)
            if not self.host:
                await self.close(code=4004)
                return

            # 验证主机状态
            if self.host.status != 'online':
                logger.warning(f"Attempt to connect to offline host {self.host.hostname}")
                await self.close(code=4003)
                return

        except Exception as e:
            logger.error(f"Failed to get host {self.host_id}: {e}")
            await self.close(code=4004)
            return

        # 接受WebSocket连接
        await self.accept()

        # 记录连接审计日志
        audit_logger.log_ssh_connection(
            user=self.user,
            host=self.host,
            action='websocket_connect',
            result='success',
            details={'client_ip': client_ip}
        )

        # 发送连接状态
        await self.send(text_data=json.dumps({
            'type': 'status',
            'message': 'WebSocket连接已建立',
            'status': 'connected'
        }))

    async def disconnect(self, close_code):
        """WebSocket断开连接"""
        await self.cleanup_ssh_connection()
        
        # 更新会话状态
        if self.session:
            await self.update_session_status('disconnected')

    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')
            
            if message_type == 'ssh_connect':
                await self.handle_ssh_connect(data)
            elif message_type == 'ssh_input':
                await self.handle_ssh_input(data)
            elif message_type == 'ssh_resize':
                await self.handle_ssh_resize(data)
            elif message_type == 'ssh_disconnect':
                await self.handle_ssh_disconnect()
            else:
                await self.send_error(f"未知的消息类型: {message_type}")
                
        except json.JSONDecodeError:
            await self.send_error("无效的JSON格式")
        except Exception as e:
            logger.error(f"Error handling message: {e}")
            await self.send_error(f"处理消息时发生错误: {str(e)}")

    async def handle_ssh_connect(self, data):
        """处理SSH连接请求"""
        try:
            # 创建SSH会话记录
            self.session = await self.create_ssh_session()
            
            # 获取连接参数
            width = data.get('width', 80)
            height = data.get('height', 24)
            
            # 建立SSH连接
            await self.establish_ssh_connection(width, height)
            
        except Exception as e:
            logger.error(f"SSH connection failed: {e}")
            await self.send_error(f"SSH连接失败: {str(e)}")
            if self.session:
                await self.update_session_status('error')

    async def handle_ssh_input(self, data):
        """处理SSH输入"""
        if not self.is_connected or not self.ssh_channel:
            await self.send_error("SSH连接未建立")
            return

        input_data = data.get('data', '')
        if input_data:
            try:
                # 检查命令安全性（如果是回车键）
                if '\r' in input_data or '\n' in input_data:
                    command = input_data.strip()
                    if command and not security_validator.is_safe_command(command):
                        logger.warning(f"Dangerous command blocked: {command} by user {self.user.username}")
                        await self.send_error("危险命令已被阻止")

                        # 记录危险命令审计日志
                        audit_logger.log_command_execution(
                            session=self.session,
                            command=command,
                            result='blocked',
                            exit_code=-1
                        )
                        return

                # 发送数据到SSH通道
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.send, input_data.encode('utf-8')
                )

                # 记录命令（如果是回车键）
                if '\r' in input_data or '\n' in input_data:
                    await self.record_command(input_data.strip())

            except Exception as e:
                logger.error(f"Failed to send SSH input: {e}")
                await self.send_error("发送数据失败")

    async def handle_ssh_resize(self, data):
        """处理终端大小调整"""
        if not self.is_connected or not self.ssh_channel:
            return
        
        width = data.get('width', 80)
        height = data.get('height', 24)
        
        try:
            await asyncio.get_event_loop().run_in_executor(
                None, self.ssh_channel.resize_pty, width, height
            )
            
            # 更新会话记录
            if self.session:
                await self.update_session_terminal_size(width, height)
                
        except Exception as e:
            logger.error(f"Failed to resize terminal: {e}")

    async def handle_ssh_disconnect(self):
        """处理SSH断开连接"""
        await self.cleanup_ssh_connection()
        await self.send(text_data=json.dumps({
            'type': 'status',
            'message': 'SSH连接已断开',
            'status': 'disconnected'
        }))

    async def establish_ssh_connection(self, width, height):
        """建立SSH连接"""
        def connect_ssh():
            try:
                # 创建SSH客户端
                self.ssh_client = paramiko.SSHClient()
                self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                
                # 连接主机
                if self.host.ssh_auth_type == 'password':
                    password = self.host.decrypt_password()
                    self.ssh_client.connect(
                        hostname=self.host.ip,
                        port=self.host.ssh_port,
                        username=self.host.ssh_username,
                        password=password,
                        timeout=10
                    )
                else:
                    private_key = self.host.decrypt_private_key()
                    if private_key:
                        from io import StringIO
                        key = paramiko.RSAKey.from_private_key(StringIO(private_key))
                        self.ssh_client.connect(
                            hostname=self.host.ip,
                            port=self.host.ssh_port,
                            username=self.host.ssh_username,
                            pkey=key,
                            timeout=10
                        )
                    else:
                        raise Exception("SSH私钥未配置")
                
                # 创建交互式shell
                self.ssh_channel = self.ssh_client.invoke_shell(
                    term='xterm',
                    width=width,
                    height=height
                )
                
                self.is_connected = True
                return True
                
            except Exception as e:
                logger.error(f"SSH connection error: {e}")
                raise e
        
        # 在线程池中执行SSH连接
        success = await asyncio.get_event_loop().run_in_executor(None, connect_ssh)
        
        if success:
            # 更新会话状态
            await self.update_session_status('connected')
            
            # 启动读取线程
            self.start_read_thread()
            
            # 发送连接成功消息
            await self.send(text_data=json.dumps({
                'type': 'status',
                'message': 'SSH连接成功',
                'status': 'connected'
            }))

    def start_read_thread(self):
        """启动SSH输出读取线程"""
        def read_ssh_output():
            while self.is_connected and self.ssh_channel:
                try:
                    if self.ssh_channel.recv_ready():
                        data = self.ssh_channel.recv(1024)
                        if data:
                            # 发送数据到WebSocket
                            asyncio.run_coroutine_threadsafe(
                                self.send_ssh_output(data.decode('utf-8', errors='ignore')),
                                asyncio.get_event_loop()
                            )
                    time.sleep(0.01)  # 避免CPU占用过高
                except Exception as e:
                    logger.error(f"Error reading SSH output: {e}")
                    break
        
        self.read_thread = threading.Thread(target=read_ssh_output, daemon=True)
        self.read_thread.start()

    async def send_ssh_output(self, data):
        """发送SSH输出到WebSocket"""
        await self.send(text_data=json.dumps({
            'type': 'ssh_output',
            'data': data
        }))

    async def cleanup_ssh_connection(self):
        """清理SSH连接"""
        self.is_connected = False
        
        if self.ssh_channel:
            try:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.close
                )
            except:
                pass
            self.ssh_channel = None
        
        if self.ssh_client:
            try:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_client.close
                )
            except:
                pass
            self.ssh_client = None

    async def send_error(self, message):
        """发送错误消息"""
        await self.send(text_data=json.dumps({
            'type': 'error',
            'message': message
        }))

    @database_sync_to_async
    def get_host(self, host_id):
        """获取主机信息"""
        try:
            return Host.objects.get(id=host_id, is_active=True)
        except Host.DoesNotExist:
            return None

    @database_sync_to_async
    def create_ssh_session(self):
        """创建SSH会话记录"""
        return SSHSession.objects.create(
            user=self.user,
            host=self.host,
            client_ip=self.get_client_ip(),
            status='connecting'
        )

    @database_sync_to_async
    def update_session_status(self, status):
        """更新会话状态"""
        if self.session:
            self.session.status = status
            if status == 'disconnected':
                self.session.end_time = timezone.now()
            self.session.save()

    @database_sync_to_async
    def update_session_terminal_size(self, width, height):
        """更新终端大小"""
        if self.session:
            self.session.terminal_width = width
            self.session.terminal_height = height
            self.session.save()

    @database_sync_to_async
    def record_command(self, command):
        """记录执行的命令"""
        if self.session and command.strip():
            SSHCommand.objects.create(
                session=self.session,
                command=command.strip()
            )

    def get_client_ip(self):
        """获取客户端IP"""
        headers = dict(self.scope.get('headers', []))
        x_forwarded_for = headers.get(b'x-forwarded-for')
        if x_forwarded_for:
            return x_forwarded_for.decode().split(',')[0].strip()
        return self.scope.get('client', ['unknown'])[0]
